# Copyright 2025 the LlamaFactory team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import TYPE_CHECKING

from ...extras import logging
from .visual import COMPOSITE_MODELS


if TYPE_CHECKING:
    from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer


logger = logging.get_logger(__name__)


def find_all_linear_modules(model: "PreTrainedModel", freeze_vision_tower: bool) -> list[str]:
    r"""Find all available modules to apply LoRA, GaLore or APOLLO."""
    # 获取模型类型（如 'llama', 'chatglm' 等）
    model_type = getattr(model.config, "model_type", None)
    # 初始化禁止插入适配器的模块集合，默认排除 lm_head 输出层
    forbidden_modules = {"lm_head"}
    # 如果是 chatglm 模型，添加 output_layer 到禁止列表
    if model_type == "chatglm":
        forbidden_modules.add("output_layer")
    # 如果是 internlm2 模型，添加 output 到禁止列表
    elif model_type == "internlm2":
        forbidden_modules.add("output")

    # 如果当前模型属于复合模型（如支持多模态），将 projector_key 加入禁止列表
    if model_type in COMPOSITE_MODELS:
        forbidden_modules.add(COMPOSITE_MODELS[model_type].projector_key)

    # 如果启用了冻结视觉塔且模型包含视觉部分，则将视觉模块加入禁止列表
    if freeze_vision_tower and model_type in COMPOSITE_MODELS:
        forbidden_modules.update(COMPOSITE_MODELS[model_type].vision_model_keys)

    # 用于保存所有符合条件的模块名称
    module_names = set()
    # 遍历模型中所有的模块（layer）
    for name, module in model.named_modules():
        # 如果模块名包含任何被禁止的关键字，跳过该模块
        if any(forbidden_module in name for forbidden_module in forbidden_modules):
            continue

        # 如果模块是 Linear 类型但不是 Embedding（避免词向量层），则将其加入候选列表
        if "Linear" in module.__class__.__name__ and "Embedding" not in module.__class__.__name__:
            # 只取模块名的最后一级（如 q_proj、v_proj 等）
            module_names.add(name.split(".")[-1])

    # 打印找到的所有线性模块名称
    logger.info_rank0("Found linear modules: {}".format(",".join(module_names)))
    # 返回模块名列表，供 LoRA 等适配器使用
    return list(module_names)


def find_expanded_modules(model: "PreTrainedModel", target_modules: list[str], num_layer_trainable: int) -> list[str]:
    r"""Find the modules in the expanded blocks to apply lora."""
    # 获取模型配置中的总层数（如 LLaMA 中的 num_hidden_layers）
    num_layers = getattr(model.config, "num_hidden_layers", None)
    # 如果没有获取到层数，说明模型不支持，抛出异常
    if not num_layers:
        raise ValueError("Model was not supported.")

    # 检查可训练层数是否能整除总层数，否则无法均匀分布
    if num_layers % num_layer_trainable != 0:
        raise ValueError(
            f"`num_layers` {num_layers} should be divisible by `num_layer_trainable` {num_layer_trainable}."
        )

    # 计算每隔多少层选择一个可训练层（stride）
    stride = num_layers // num_layer_trainable
    # 根据 stride 确定哪些层 ID 是可训练的（从第 stride - 1 层开始选）
    trainable_layer_ids = range(stride - 1, num_layers + stride - 1, stride)
    # 构建可训练层名格式，如 ".12.", ".24." 等
    trainable_layers = [f".{idx:d}." for idx in trainable_layer_ids]
    # 存储匹配的模块名称
    module_names = []
    # 遍历模型所有模块
    for name, _ in model.named_modules():
        # 同时满足：
        # 1. 名称包含目标模块（如 q_proj、v_proj 等）
        # 2. 所属层是被选中的可训练层
        if any(target_module in name for target_module in target_modules) and any(
            trainable_layer in name for trainable_layer in trainable_layers
        ):
            module_names.append(name)

    # 打印应用 LoRA 的具体层编号
    logger.info_rank0("Apply lora to layers: {}.".format(",".join(map(str, trainable_layer_ids))))
    # 返回最终匹配的模块名称列表
    return module_names


def register_autoclass(config: "PretrainedConfig", model: "PreTrainedModel", tokenizer: "PreTrainedTokenizer"):
    # 如果配置中包含 AutoConfig 的自动映射，则注册该类为 auto class
    if "AutoConfig" in getattr(config, "auto_map", {}):
        config.__class__.register_for_auto_class()
    # 如果配置中包含 AutoModelForCausalLM 的自动映射，则注册模型类为 auto class
    if "AutoModelForCausalLM" in getattr(config, "auto_map", {}):
        model.__class__.register_for_auto_class()
    # 如果 tokenizer 的 init_kwargs 中有 AutoTokenizer 的自动映射，则注册 tokenizer 类为 auto class
    if "AutoTokenizer" in tokenizer.init_kwargs.get("auto_map", {}):
        tokenizer.__class__.register_for_auto_class()
